import os
import sys
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
import pandas as pd
import tushare as ts
from utils.confutils import ConfUtil
from tqdm import tqdm
import time

# 提取token
token = ConfUtil().get_tokens('tushare')
pro = ts.pro_api(token)
# 加载申万分类信息
swclass_file = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) + "\\Data\\RawData\\index_sw_class.feather"
df_swclass = pd.read_feather(swclass_file).query("level=='L3'")

df_full = []
print('开始下载指数成分股信息...')
for indexcode in tqdm(df_swclass.index_code.unique()):
    for is_new in ['Y','N']:
        while 1==1:
            try:
                df = pro.index_member_all(l3_code=indexcode,is_new=is_new)
                break
            except Exception as e:
                time.sleep(5)
                continue
        df_full.append(df)
df_res = pd.concat(df_full,axis=0)
print('下载完毕')
print(f'数据信息：\n{df_res.info()}')
print(df_res.head(5))
# 存到.\Data\RawData\index_sw_member.feather文件中
tofile = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) + "\\Data\\RawData\\index_sw_member.feather"
df_res.to_feather(tofile)